# coding=utf-8
# @author:      ChengJing
# @name:        analyze_data.py
# @datetime:    2021/12/11 18:05
# @software:    PyCharm
# @description:

import numpy as np
import pandas as pd
import wntr
import re
import os


class AnalyzeData:
    """
    分析爆管数据：
    --变量：a、爆管程度
    --对于每一个节点：
          1. 爆管点引起的监测点集合的压降（最小、中位数、平均数、最大）
          2. 爆管点的爆管流量（最小、中位数、平均数、最大）
          3. 节点的管径
    Return:
        junction | diameter | burst_q | delta_pressure_min | delta_pressure_media | delta_pressure_mean | delta_pressure_max
    """
    def __init__(self, root, save_file):
        self.wn = wntr.network.WaterNetworkModel(r'./inp/tmodel24.inp')
        self.junctions = self.wn.junction_name_list
        self.root = root
        self.save_file = root + '/' + save_file

    def index2id(self, indexs):
        ids = [self.junctions[i] for i in indexs]
        return ids

    def _get_datas(self, junction):
        files = os.listdir(self.root)
        data = []
        burst_level = []
        for f in files:
            match = re.match(r'P_'+junction+r'_'+r'([0-9.]*)\.csv', f)
            if match:
                delta_p, index, _ = self._analyze_pressure(self.root+'/'+match.group(0), junction)
                delta_q = self._analyze_flow(self.root+'/'+match.group(0).replace('P', 'Q'), index)
                pq = np.zeros((9,))
                pq[1::2] = delta_q
                pq[2::2] = delta_p
                pq[0] = max([self.wn.get_link(i).diameter for i in self.wn.get_links_for_node(junction)]) * 1000
                pq = list(pq)
                pq.insert(0, junction)
                data.append(pq)
                burst_level.append(float(match.group(1)))
        return data, burst_level

    def _get_monitors_datas(self, junction, monitors):
        files = os.listdir(self.root)
        data = []
        burst_level = []
        for f in files:
            match = re.match(r'P_'+junction+r'_'+r'([0-9.]*)\.csv', f)
            if match:
                pq = np.zeros((len(monitors)+2,))
                delta_p, index, df = self._analyze_pressure(self.root+'/'+match.group(0), junction)
                delta_q = self._analyze_flow(self.root+'/'+match.group(0).replace('P', 'Q'), index)
                pq[0] = delta_q[-1]
                pq[1] = delta_p[-1]
                for i, n in enumerate(monitors,start=2):
                    delta_p, index = self._analyze_pressure2(df, n)
                    pq[i] = delta_p[-1]
                pq = list(pq)
                pq.insert(0, junction)
                data.append(pq)
                burst_level.append(float(match.group(1)))
        return data, burst_level

    def _analyze_pressure(self, file, junction):
        df = pd.read_csv(file, header=0, index_col=0)
        data = df[junction].values
        delta = data[60::60] - data[30:-30:60]
        index = delta.argsort()[[0, int(len(delta)/2), -1]]
        return list(delta[index])+[delta.mean()], index, df

    def _analyze_pressure2(self, df, junction):
        data = df[junction].values
        delta = data[60::60] - data[30:-30:60]
        index = delta.argsort()[[0, int(len(delta)/2), -1]]
        return list(delta[index])+[delta.mean()], index

    def _analyze_flow(self, file, index):
        df = pd.read_csv(file, header=0, index_col=0)
        data = df.values
        delta = data[30:-30:60]
        return list(delta[index])+[delta.mean()]

    def _analyze_junction(self, data):
        data = np.array(data)
        pass

    def analyze(self, burst_level_len):
        writer = pd.ExcelWriter(self.save_file)
        df = []
        sheet_name = []
        for i in range(burst_level_len):
            df.append([])
        for junction in self.junctions:
            data, burst_level = self._get_datas(junction)
            # self._analyze_junction(data)
            if data:
                if not sheet_name:
                    sheet_name = burst_level
                for m, n in enumerate(burst_level):
                    df[m].append(data[m])
        for i in range(burst_level_len):
            dataframe = pd.DataFrame(df[i], columns=[
                'junction', 'diameter',
                'burst_flow_min(m3/s)', 'burst_pressure_drop_min(m)',
                'burst_flow_median(m3/s)', 'burst_pressure_drop_median(m)',
                'burst_flow_max(m3/s)', 'burst_pressure_drop_max(m)',
                'burst_flow_mean(m3/s)', 'burst_pressure_drop_mean(m)'
            ])
            dataframe.to_excel(writer, sheet_name=str(sheet_name[i]))
        writer.close()

    def analyze_monitor(self, burst_level_len, monitors):
        writer = pd.ExcelWriter(self.save_file)
        df = []
        sheet_name = []
        for i in range(burst_level_len):
            df.append([])
        for junction in self.junctions:
            data, burst_level = self._get_monitors_datas(junction, monitors)
            # self._analyze_junction(data)
            if data:
                if not sheet_name:
                    sheet_name = burst_level
                for m, n in enumerate(burst_level):
                    df[m].append(data[m])
        for i in range(burst_level_len):
            dataframe = pd.DataFrame(df[i], columns=[
                'junction', 'burst_flow_mean(m3/s)', 'burst_pressure_drop_mean(m)'] + [f'{i}_pressure_drop_mean(m)' for i in monitors]
                )
            dataframe.to_excel(writer, sheet_name=str(sheet_name[i]))
        writer.close()


if __name__ == '__main__':
    monitors = [456, 349, 220, 130, 70, 422, 405, 402, 452, 30, 272, 150, 156]
    model = AnalyzeData(r'F:\模型数据\new_data\burst', '分析监测点1.xlsx')
    ids = model.index2id(monitors)
    model.analyze_monitor(6, ids)
